"""
ChatTTS 音色生成和测试工具

功能：
1. 生成随机音色并试听
2. 保存喜欢的音色
3. 测试已保存的音色
"""

import os
import sys
import torch
import ChatTTS
import torchaudio
import numpy as np
from pathlib import Path


class SpeakerGenerator:
    """音色生成器"""
    
    def __init__(self):
        """初始化ChatTTS"""
        print("=" * 70)
        print("ChatTTS 音色生成和测试工具")
        print("=" * 70)
        
        print("\n[1/2] 正在初始化ChatTTS引擎...")
        self.chat = ChatTTS.Chat()
        self.chat.load(compile=False, source='huggingface')
        print("✓ ChatTTS引擎初始化完成")
        
        # 创建音色保存目录
        self.speaker_dir = Path("speakers")
        self.speaker_dir.mkdir(exist_ok=True)
        print(f"✓ 音色保存目录: {self.speaker_dir.absolute()}")
    
    def generate_random_speakers(self, count: int = 5, test_text: str = "大家好,我是AI虚拟主播小爱!"):
        """生成多个随机音色并试听
        
        Args:
            count: 生成数量
            test_text: 测试文本
        """
        print(f"\n[2/2] 正在生成 {count} 个随机音色...")
        print(f"测试文本: {test_text}")
        print("-" * 70)
        
        speakers = []
        
        for i in range(count):
            print(f"\n音色 #{i+1}/{count}")
            
            # 生成随机音色
            spk_emb = self.chat.sample_random_speaker()
            speakers.append(spk_emb)
            
            # 合成测试音频
            params_infer_code = ChatTTS.Chat.InferCodeParams(
                spk_emb=spk_emb,
                temperature=0.3,
                top_P=0.7,
                top_K=20,
            )
            
            wavs = self.chat.infer([test_text], params_infer_code=params_infer_code)
            
            if wavs is not None and len(wavs) > 0:
                # 保存试听文件
                output_file = f"speaker_preview_{i+1}.wav"
                self._save_wav(output_file, wavs[0])
                print(f"  ✓ 试听文件: {output_file}")
            else:
                print(f"  ✗ 生成失败")
        
        print("\n" + "=" * 70)
        print("音色生成完成!")
        print(f"请试听 speaker_preview_1.wav 到 speaker_preview_{count}.wav")
        print("=" * 70)
        
        return speakers
    
    def save_speaker(self, speaker_index: int, speakers: list, name: str = ""):
        """保存指定的音色
        
        Args:
            speaker_index: 音色序号(1-based)
            speakers: 音色列表
            name: 保存的文件名(不含扩展名)
        """
        if speaker_index < 1 or speaker_index > len(speakers):
            print(f"✗ 无效的音色序号: {speaker_index}")
            return None
        
        spk_emb = speakers[speaker_index - 1]
        
        if name is None or name == "":
            name = f"speaker_{speaker_index}"
        
        filepath = self.speaker_dir / f"{name}.pt"
        torch.save(spk_emb, filepath)
        print(f"✓ 音色已保存: {filepath}")
        return filepath
    
    def generate_with_seed(self, seed: int, test_text: str = "大家好,我是AI虚拟主播小爱!"):
        """使用固定种子生成音色
        
        Args:
            seed: 随机种子
            test_text: 测试文本
        """
        print(f"\n使用种子 {seed} 生成音色...")
        
        # 设置随机种子
        torch.manual_seed(seed)
        spk_emb = self.chat.sample_random_speaker()
        
        # 合成测试音频
        params_infer_code = ChatTTS.Chat.InferCodeParams(
            spk_emb=spk_emb,
            temperature=0.3,
            top_P=0.7,
            top_K=20,
        )
        
        wavs = self.chat.infer([test_text], params_infer_code=params_infer_code)
        
        if wavs is not None and len(wavs) > 0:
            output_file = f"speaker_seed_{seed}.wav"
            self._save_wav(output_file, wavs[0])
            print(f"✓ 试听文件: {output_file}")
        
        # 保存音色
        filepath = self.speaker_dir / f"seed_{seed}.pt"
        torch.save(spk_emb, filepath)
        print(f"✓ 音色已保存: {filepath}")
        
        return spk_emb
    
    def test_speaker_file(self, filepath: str, test_text: str = "大家好,我是AI虚拟主播小爱!"):
        """测试已保存的音色文件
        
        Args:
            filepath: 音色文件路径(.pt格式)
            test_text: 测试文本
        """
        print(f"\n正在测试音色文件: {filepath}")
        
        if not os.path.exists(filepath):
            print(f"✗ 文件不存在: {filepath}")
            return
        
        # 加载音色
        spk_emb = torch.load(filepath, weights_only=True)
        print(f"✓ 音色加载成功")
        
        # 合成测试音频
        params_infer_code = ChatTTS.Chat.InferCodeParams(
            spk_emb=spk_emb,
            temperature=0.3,
            top_P=0.7,
            top_K=20,
        )
        
        wavs = self.chat.infer([test_text], params_infer_code=params_infer_code)
        
        if wavs is not None and len(wavs) > 0:
            output_file = f"test_{Path(filepath).stem}.wav"
            self._save_wav(output_file, wavs[0])
            print(f"✓ 试听文件: {output_file}")
        else:
            print(f"✗ 生成失败")
    
    def _save_wav(self, filepath: str, audio_data: np.ndarray, sample_rate: int = 24000):
        """保存为WAV文件"""
        try:
            audio_tensor = torch.from_numpy(audio_data).unsqueeze(0)
            torchaudio.save(filepath, audio_tensor, sample_rate)
        except Exception as e:
            # 备用方法
            import wave
            audio_int16 = (audio_data * 32767).astype(np.int16)
            with wave.open(filepath, 'wb') as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)
                wf.setframerate(sample_rate)
                wf.writeframes(audio_int16.tobytes())


def main():
    """主函数"""
    generator = SpeakerGenerator()
    
    print("\n" + "=" * 70)
    print("请选择操作:")
    print("  1. 生成5个随机音色并试听")
    print("  2. 使用固定种子生成音色(推荐女声种子: 2222, 6666, 9527)")
    print("  3. 测试已保存的音色文件")
    print("=" * 70)
    
    choice = input("\n请输入选项(1/2/3): ").strip()
    
    if choice == "1":
        # 生成随机音色
        speakers = generator.generate_random_speakers(5)
        
        # 询问是否保存
        print("\n如果你喜欢某个音色,可以保存它")
        save_choice = input("请输入要保存的音色序号(1-5),或按Enter跳过: ").strip()
        
        if save_choice.isdigit():
            index = int(save_choice)
            name = input("请输入音色名称(如: female_1): ").strip()
            generator.save_speaker(index, speakers, name)
    
    elif choice == "2":
        # 使用固定种子
        print("\n推荐的女声种子:")
        print("  2222 - 温柔女声")
        print("  6666 - 活泼女声")
        print("  9527 - 甜美女声")
        print("  1234 - 清新女声")
        
        seed_input = input("\n请输入种子(或使用推荐值): ").strip()
        seed = int(seed_input) if seed_input.isdigit() else 2222
        
        generator.generate_with_seed(seed)
    
    elif choice == "3":
        # 测试音色文件
        filepath = input("请输入音色文件路径(.pt文件): ").strip()
        generator.test_speaker_file(filepath)
    
    else:
        print("无效的选项")


if __name__ == "__main__":
    main()
